from langchain.llms.base import LLM
from transformers import AutoTokenizer, AutoModel
from langchain_core.messages.ai import AIMessage
from typing import List

class ChatGLM3(LLM):
    # 类属性设置
    max_token: int = 8192  # 最大生成的token数量
    do_sample: bool = True  # 是否进行随机采样
    temperature: float = 0.3  # 温度控制生成的随机性
    top_p: float = 0.0  # 采样的top-p参数
    tokenizer: AutoTokenizer = None  # 存储分词器对象
    model: AutoModel = None  # 存储模型对象
    history: List[dict] = []  # 存储对话历史

    def __init__(self):
        super().__init__()  # 初始化父类

    @property
    def _llm_type(self):
        # 返回模型类型
        return "ChatGLM3"

    def load_model(self, modelPath: str):
        """
        加载模型和分词器
        :param modelPath: 模型路径
        """
        # 配置分词器
        self.tokenizer = AutoTokenizer.from_pretrained(modelPath, trust_remote_code=True, use_fast=True)

        # 加载模型
        self.model = AutoModel.from_pretrained(modelPath, trust_remote_code=True, device_map="auto")

        # 设置模型为评估模式
        self.model.eval()

    def _call(self, prompt: str, config: dict = {}, history: List[dict] = None) -> AIMessage:
        """
        调用模型生成响应
        :param prompt: 输入的提示语
        :param config: 配置参数
        :param history: 对话历史
        :return: AIMessage对象
        """
        if history is None:
            history = []
        return self.invoke(prompt, history)

    def invoke(self, prompt: str, config: dict = {}, history: List[dict] = None) -> AIMessage:
        """
        生成响应的核心逻辑
        :param prompt: 输入的提示语
        :param config: 配置参数
        :param history: 对话历史
        :return: AIMessage对象
        """
        if not isinstance(prompt, str):
            prompt = prompt.to_string()  # 确保输入是字符串
        # 使用模型生成响应
        response, history = self.model.chat(
            self.tokenizer,
            prompt,
            history=history,
            do_sample=self.do_sample,
            max_length=self.max_token,
            temperature=self.temperature
        )
        self.history = history  # 更新对话历史
        return AIMessage(content=response)  # 返回生成的消息

    def stream(self, prompt: str, config: dict = {}, history: List[dict] = None):
        """
        流式生成响应
        :param prompt: 输入的提示语
        :param config: 配置参数
        :param history: 对话历史
        """
        if history is None:
            history = []
        if not isinstance(prompt, str):
            prompt = prompt.to_string()  # 确保输入是字符串
        preResponse = ""  # 上一个响应的内容
        # 使用流式聊天生成响应
        for response, new_history in self.model.stream_chat(self.tokenizer, prompt):
            # 获取当前响应与上一个响应的差异
            if preResponse == "":
                result = response  # 第一个响应
            else:
                result = response[len(preResponse):]  # 新增部分
            preResponse = response  # 更新上一个响应
            yield result  # 逐步返回生成的内容
